#pragma once
#include"Common.h"
#include "VTCore.h"

/// <summary>
/// PML4, PDPTE, PDTE pointing to another table
/// </summary>
typedef union _EPT_MMPTE
{
	ULONG64 All;
	struct
	{
		ULONG64 Present : 1;    // If the region is present (read access)
		ULONG64 Write : 1;      // If the region is writable
		ULONG64 Execute : 1;    // If the region is executable
		ULONG64 reserved1 : 9;  // Reserved
		ULONG64 PhysAddr : 40;  // Physical address
		ULONG64 reserved2 : 12; // Reserved
	} Fields;
} EPT_PML4_ENTRY, EPT_MMPTE, *PEPT_PML4_ENTRY, *PEPT_MMPTE;


/// <summary>
/// Exit qualification for EPT violation
/// </summary>
typedef union _EPT_VIOLATION_DATA
{
	ULONG64 All;
	struct
	{
		ULONG64 Read : 1;           // Read access
		ULONG64 Write : 1;          // Write access
		ULONG64 Execute : 1;        // Execute access
		ULONG64 PTERead : 1;        // PTE entry has read access
		ULONG64 PTEWrite : 1;       // PTE entry has write access
		ULONG64 PTEExecute : 1;     // PTE entry has execute access
		ULONG64 Reserved1 : 1;      // 
		ULONG64 GuestLinear : 1;    // GUEST_LINEAR_ADDRESS field is valid
		ULONG64 FailType : 1;       // 
		ULONG64 Reserved2 : 3;      // 
		ULONG64 NMIBlock : 1;       // NMI unblocking due to IRET
		ULONG64 Reserved3 : 51;     // 
	} Fields;
} EPT_VIOLATION_DATA, *PEPT_VIOLATION_DATA;

/// <summary>
/// ETP table level
/// </summary>
typedef enum _EPT_TABLE_LEVEL
{
	EPT_LEVEL_PTE = 0,
	EPT_LEVEL_PDE = 1,
	EPT_LEVEL_PDPTE = 2,
	EPT_LEVEL_PML4 = 3,
	EPT_TOP_LEVEL = EPT_LEVEL_PML4
} EPT_TABLE_LEVEL;



/// <summary>
/// EPT page access
/// </summary>
typedef enum _EPT_ACCESS
{
	EPT_ACCESS_NONE = 0,
	EPT_ACCESS_READ = 1,
	EPT_ACCESS_WRITE = 2,
	EPT_ACCESS_EXEC = 4,
	EPT_ACCESS_RW = EPT_ACCESS_READ | EPT_ACCESS_WRITE,
	EPT_ACCESS_ALL = EPT_ACCESS_READ | EPT_ACCESS_WRITE | EPT_ACCESS_EXEC
} EPT_ACCESS;


/// <summary>
/// PTE entry
/// </summary>
typedef union _EPT_PTE_ENTRY
{
	ULONG64 All;
	struct
	{
		ULONG64 Read : 1;           // Region is present (read access)
		ULONG64 Write : 1;          // Region is writable
		ULONG64 Execute : 1;        // Region is executable
		ULONG64 MemoryType : 3;     // EPT Memory type
		ULONG64 IgnorePat : 1;      // Flag for whether to ignore PAT
		ULONG64 reserved1 : 5;      // Reserved
		ULONG64 PhysAddr : 40;      // Physical address
		ULONG64 reserved2 : 12;     // Reserved
	} Fields;
} EPT_PTE_ENTRY, *PEPT_PTE_ENTRY;




typedef union _VMX_CPU_BASED_CONTROLS
{
	ULONG32 All;
	struct
	{
		ULONG32 Reserved1 : 2;                 // [0-1]
		ULONG32 InterruptWindowExiting : 1;    // [2]
		ULONG32 UseTSCOffseting : 1;           // [3]
		ULONG32 Reserved2 : 3;                 // [4-6]
		ULONG32 HLTExiting : 1;                // [7]
		ULONG32 Reserved3 : 1;                 // [8]
		ULONG32 INVLPGExiting : 1;             // [9]
		ULONG32 MWAITExiting : 1;              // [10]
		ULONG32 RDPMCExiting : 1;              // [11]
		ULONG32 RDTSCExiting : 1;              // [12]
		ULONG32 Reserved4 : 2;                 // [13-14]
		ULONG32 CR3LoadExiting : 1;            // [15]
		ULONG32 CR3StoreExiting : 1;           // [16]
		ULONG32 Reserved5 : 2;                 // [17-18]
		ULONG32 CR8LoadExiting : 1;            // [19]
		ULONG32 CR8StoreExiting : 1;           // [20]
		ULONG32 UseTPRShadowExiting : 1;       // [21]
		ULONG32 NMIWindowExiting : 1;          // [22]
		ULONG32 MovDRExiting : 1;              // [23]
		ULONG32 UnconditionalIOExiting : 1;    // [24]
		ULONG32 UseIOBitmaps : 1;              // [25]
		ULONG32 Reserved6 : 1;                 // [26]
		ULONG32 MonitorTrapFlag : 1;           // [27]
		ULONG32 UseMSRBitmaps : 1;             // [28]
		ULONG32 MONITORExiting : 1;            // [29]
		ULONG32 PAUSEExiting : 1;              // [30]
		ULONG32 ActivateSecondaryControl : 1;  // [31]
	} Fields;
} VMX_CPU_BASED_CONTROLS, *PVMX_CPU_BASED_CONTROLS;

typedef union _VMX_SECONDARY_CPU_BASED_CONTROLS
{
	ULONG32 All;
	struct
	{
		ULONG32 VirtualizeAPICAccesses : 1;      // [0]
		ULONG32 EnableEPT : 1;                   // [1]
		ULONG32 DescriptorTableExiting : 1;      // [2]
		ULONG32 EnableRDTSCP : 1;                // [3]
		ULONG32 VirtualizeX2APICMode : 1;        // [4]
		ULONG32 EnableVPID : 1;                  // [5]
		ULONG32 WBINVDExiting : 1;               // [6]
		ULONG32 UnrestrictedGuest : 1;           // [7]
		ULONG32 APICRegisterVirtualization : 1;  // [8]
		ULONG32 VirtualInterruptDelivery : 1;    // [9]
		ULONG32 PAUSELoopExiting : 1;            // [10]
		ULONG32 RDRANDExiting : 1;               // [11]
		ULONG32 EnableINVPCID : 1;               // [12]
		ULONG32 EnableVMFunctions : 1;           // [13]
		ULONG32 VMCSShadowing : 1;               // [14]
		ULONG32 Reserved1 : 1;                   // [15]
		ULONG32 RDSEEDExiting : 1;               // [16]
		ULONG32 Reserved2 : 1;                   // [17]
		ULONG32 EPTViolation : 1;                // [18]
		ULONG32 Reserved3 : 1;                   // [19]
		ULONG32 EnableXSAVESXSTORS : 1;          // [20]
	} Fields;
} VMX_SECONDARY_CPU_BASED_CONTROLS, *PVMX_SECONDARY_CPU_BASED_CONTROLS;


#pragma pack(push, 1)

/// <summary>
///  PEPT
/// </summary>
typedef union _EPT_TABLE_POINTER
{
	ULONG64 All;
	struct
	{
		ULONG64 MemoryType : 3;         // EPT Paging structure memory type (0 for UC)
		ULONG64 PageWalkLength : 3;     // Page-walk length
		ULONG64 reserved1 : 6;
		ULONG64 PhysAddr : 40;          // Physical address of the EPT PML4 table
		ULONG64 reserved2 : 12;
	} Fields;
} EPT_TABLE_POINTER, *PEPT_TABLE_POINTER;

#pragma pack(pop)



NTSTATUS EptBuildIdentityMap(IN PEPT_DATA pEPT);
VOID EptEnable(IN PEPT_PML4_ENTRY PML4);
VOID VmExitEptViolation();


NTSTATUS EptUpdateTableRecursive(
	IN PEPT_DATA pEPTData,
	IN PEPT_MMPTE pTable,
	IN EPT_TABLE_LEVEL level,
	IN ULONG64 pfn,
	IN UCHAR access,
	IN ULONG64 hostPFN,
	IN ULONG count
);

